from core.cuda import cuda_module
from core.function import Function
from implement.functions.conv.pooling2d_withindexes import Pooling2DWithIndexes
from utils.common import pair
from utils.functions_conv import col2im_array


class Pooling2DGrad(Function):
    def __init__(self, mpool2d):
        """
        池化操作的反向传播函数类，用于计算输入变量的梯度。

        Args:
            mpool2d (Pooling): 对应的前向传播池化操作函数。
        """
        self.mpool2d = mpool2d
        self.kernel_size = mpool2d.kernel_size
        self.stride = mpool2d.stride
        self.pad = mpool2d.pad
        self.input_shape = mpool2d.inputs[0].shape
        self.dtype = mpool2d.inputs[0].dtype
        self.indexes = mpool2d.indexes

    def forward(self, gy):
        """
        前向传播，计算输入变量的梯度。

        Args:
            gy (Variable): 输出变量的梯度。

        Returns:
            Variable: 输入变量的梯度。
        """
        xp = cuda_module

        N, C, OH, OW = gy.shape
        N, C, H, W = self.input_shape
        KH, KW = pair(self.kernel_size)

        gcol = xp.zeros((N * C * OH * OW * KH * KW), dtype=self.dtype)

        indexes = (self.indexes.ravel()
                   + xp.arange(0, self.indexes.size * KH * KW, KH * KW))

        gcol[indexes] = gy.ravel()
        gcol = gcol.reshape(N, C, OH, OW, KH, KW)
        gcol = xp.swapaxes(gcol, 2, 4)
        gcol = xp.swapaxes(gcol, 3, 5)

        gx = col2im_array(gcol, (N, C, H, W), self.kernel_size, self.stride,
                          self.pad, to_matrix=False)
        return gx

    def backward(self, ggx):
        """
        反向传播，计算输入变量的二阶梯度。

        Args:
            ggx (Variable): 输出变量的梯度。

        Returns:
            Variable: 输入变量的二阶梯度。
        """
        f = Pooling2DWithIndexes(self.mpool2d)
        return f(ggx)

